# coding=utf-8
from engine.data.dataloader import D
from engine.performance import PerformanceUtils
import pandas as pd
from utils.mongo_utils import get_db
import pymongo
import matplotlib.pyplot as plt
import talib as ta

def analysis_code(code):
    df = D.load([code],fields=['$close/ Ref($close,20) - 1','MACD($close)','$close/Ref($close,1) -1'], names=['mom_20','macd','rate'])
    df['equity'] = (df['rate']+1).cumprod()
    ret = PerformanceUtils().calc_equity(df[['equity']])
    df.index = pd.to_datetime(df['date'])
    #print(ret)

    macd_day = round(df['macd'][-1], 2)
    mom_20 = round(df['mom_20'][-1], 2)
    print(f'{code}的20日动量值为{mom_20}, macd值={macd_day}'.format(code=code,mom_20=mom_20,macd=macd_day))

    #df['equity'].plot()

    df_w = df[['close']].resample('w').last()
    df_w['high'] = df['high'].resample('w').max()
    df_w['low'] = df['low'].resample('w').min()

    #print(df_w)
    import talib as ta

    df_w.dropna(inplace=True)
    macd,_,_ = ta.MACD(df_w['close'].values, 12, 26, 9)

    df['slowk'], df['slowd'] = ta.STOCH(
        df['high'].values,
        df['low'].values,
        df['close'].values,
        fastk_period=9,
        slowk_period=3,
        slowk_matype=0,
        slowd_period=3,
        slowd_matype=0)
    # 求出J值，J = (3*K)-(2*D)
    df['slowj'] = list(map(lambda x, y: 3 * x - 2 * y, df['slowk'], df['slowd']))


    print('周线MACD:',macd[-3],macd[-2],macd[-1], )
    kdj = df['slowj'][-1]
    status = '超买' if kdj > 50 else '超卖'
    print('日线KDJ',kdj,status, '2日劲道：', (df['close'][-1] - df['close'][-3])/2)
    #print(df_w)
    #plt.show()

def query_funds():
    f = open('mgr.txt', 'r')
    lines = f.readlines()
    print(lines)
    codes = [line.replace('\n','') for line in lines]
    print(codes)

    query = {}
    query['ts_code'] = {'$in': codes}

    '''
    #激进债
    query['1y_annu_return']={'$gt':0.05}
    query['2y_annu_return'] = {'$gt': 0.05}
    query['3y_annu_return']={'$gte':0.1}
    query['3y_sharpe'] = {'$gte': 1}
    query['3y_mdd']= {'$gte': -0.1, '$lte': -0.01}
    '''

    query['1y_annu_return'] = {'$gt': 0.15}
    query['2y_annu_return'] = {'$gt': 0.15}
    query['3y_annu_return'] = {'$gte': 0.15}

    query['3y_mdd'] = {'$gte': -0.3, '$lte': -0.01}


    items = get_db()['factors_funds'].find(query,{'3y_mdd':1,'3y_annu_return':1,'ts_code':1,'name':1,'3y_sharpe':1,'3y_calmar':1})
    df = pd.DataFrame(list(items))

    df['sorter'] = 0.5 * df['3y_sharpe'] + 0.5 * df['3y_calmar']
    #print(df)
    df.sort_values(by='sorter',ascending=False,inplace=True)

    print(df)
    df.to_csv('3y.csv')

def plot_code(code):
    df = D.load([code])
    df['rate'] = df['adj_nav'].pct_change()
    df['equity'] = (df['rate'] + 1).cumprod()
    ret = PerformanceUtils().calc_equity(df[['equity']])
    df.index = pd.to_datetime(df['date'])
    print(ret)
    df['equity'].plot()
    import matplotlib.pyplot as plt

    plt.show()

if __name__ == '__main__':
    for code in ['000300.SH','000905.SH','000852.SH', '399006.SZ','399324.SZ','399997.SZ','000013.SH','HSI','N225','SPX','GDAXI']:
        analysis_code(code)
    #query_funds()
    #plot_code('002363.OF')

